from typing import Any, Optional, Dict, List, Tuple
import calendar
import re
from datetime import datetime, timedelta

from MCPService import amap_mcp_server
from core.config import settings

from langchain.tools import Tool
from langchain_core.tools import tool

from core import OllamaEmbeddings_Create
from core.Milvus import GET_Milvus  # 确保GET_Milvus是单例模式
from langchain_core.documents import Document

# ------------------------------
# 新增：导入智谱官方SDK（需提前安装：pip install zai-sdk）
# ------------------------------
from zai import ZhipuAiClient


# ------------------------------
# 1. 知识库核心处理类（保持原有逻辑）
# ------------------------------
class Knowledge:
    def __init__(self):
        """初始化知识库（依赖GET_Milvus单例的固定集合配置）"""
        self.milvus = GET_Milvus()
        self.vector_store = self.milvus.init_fixed_collection()
        self.fixed_collection = self.milvus.FIXED_COLLECTION_NAME
        self.fixed_dim = self.milvus.FIXED_VECTOR_DIM

    def add_documents(self,
                      documents: List[Document],
                      ids: Optional[List[str]] = None) -> Tuple[bool, List[str]]:
        if not documents:
            print("❌ 没有文档可添加")
            return False, []
        try:
            if not self.vector_store:
                print("❌ 向量存储未初始化，添加失败")
                return False, []
            inserted_ids = self.vector_store.add_documents(documents=documents, ids=ids if ids else None)
            print(f"✅ 成功添加 {len(inserted_ids)} 个文档到集合[{self.fixed_collection}]")
            return True, inserted_ids
        except Exception as e:
            print(f"❌ 添加文档失败: {str(e)}")
            return False, []

    def add_texts(self,
                  texts: List[str],
                  metadatas: Optional[List[Dict[str, Any]]] = None,
                  ids: Optional[List[str]] = None) -> Tuple[bool, List[str]]:
        if not texts:
            print("❌ 没有文本可添加")
            return False, []
        documents = []
        for i, text in enumerate(texts):
            metadata = metadatas[i] if (metadatas and i < len(metadatas)) else {}
            documents.append(Document(page_content=text.strip(), metadata=metadata))
        return self.add_documents(documents, ids)

    def search_similar(self,
                       query: str,
                       k: int = 5,
                       filter: Optional[Dict[str, Any]] = None) -> List[Document]:
        # try:
            if not self.vector_store:
                print("❌ 向量存储未初始化，搜索失败")
                return []

            # 构建基础搜索参数，**不包含output_fields**
            search_kwargs = {
                "k": min(k, 20),
                "fetch_k": min(k * 3, 40),
                "lambda_mult": 0.5,

            }
            # 仅当filter不为None时，才添加filter参数
            if filter is not None:
                search_kwargs["filter"] = filter

            mmr_retriever = self.vector_store.as_retriever(
                search_type="mmr",
                search_kwargs=search_kwargs
            )

            docs = mmr_retriever.invoke(query.strip())
            print(f"✅ 在集合[{self.fixed_collection}]中找到 {len(docs)} 个相似文档（MMR算法）")
            print(docs)
            return docs
        # except Exception as e:
        #     print(f"❌ MMR搜索失败: {str(e)}")
        #     return []

    def delete_document(self, document_id: str) -> bool:
        try:
            if not self.vector_store:
                print("❌ 向量存储未初始化，删除失败")
                return False
            self.vector_store.delete(ids=[str(document_id)])
            print(f"✅ 文档ID[{document_id}]已从集合[{self.fixed_collection}]删除")
            return True
        except Exception as e:
            print(f"❌ 删除文档失败: {str(e)}")
            return False

    def clear_collection(self) -> bool:
        try:
            self.milvus.client.delete(collection_name=self.fixed_collection, expr="id >= 0")
            print(f"✅ 集合[{self.fixed_collection}]已清空所有文档")
            return True
        except Exception as e:
            print(f"❌ 清空集合失败: {str(e)}")
            return False

    def get_document_count(self) -> int:
        try:
            stats = self.milvus.client.get_collection_stats(self.fixed_collection)
            count = stats.get("row_count", 0)
            print(f"✅ 集合[{self.fixed_collection}]当前文档数: {count}")
            return count
        except Exception as e:
            print(f"❌ 获取文档数量失败: {str(e)}")
            return 0


# ------------------------------
# 2. 日期解析工具（修复节日日期解包错误）
# ------------------------------
@tool("date_parser", return_direct=False)
def date_parser(query: str) -> str:
    """
    【适用场景】
    1. 解析**相对日期**（如"大前天"、"大后天"、"上周三"）；
    2. 查询**绝对日期**（如"2025年中秋节"、"国庆节"）；
    3. 获取日期的**星期信息**（如"2025年10月1日是星期几"）。

    【输入格式】
    - 直接传入自然语言日期描述，例如：
      "大前天的日期"
      "2025年中秋节是几号"
      "2025年10月1日星期几"

    【输出说明】
    - 格式：YYYY-MM-DD 星期X
    - 支持的相对偏移量：-4天（大前天）至+4天（大后天）
    - 支持的节日：春节、元宵节、清明节、端午节、中秋节、国庆节

    【示例】
    - 输入："大后天"
    - 输入："今年的中秋节"
    - 输入："2025年国庆节"
    """
    try:
        # 处理相对日期
        relative_map = {
            "今天": 0, "明天": 1, "后天": 2, "大后天": 3,
            "昨天": -1, "前天": -2, "大前天": -3, "大大前天": -4
        }
        if query in relative_map:
            delta = relative_map[query]
            target_date = datetime.now() + timedelta(days=delta)
            return f"{target_date.strftime('%Y-%m-%d')} {calendar.day_name[target_date.weekday()]}"

        # 处理绝对日期查询（如"2025年中秋节"）
        query_lower = query.lower()

        # 预定义节日日期（修复腊八节解包问题：统一用(year, month, day, desc)格式）
        holidays = {
            # 传统节日（2025年公历对应日期）
            "春节": (2025, 1, 29, "农历正月初一，中国最重要的传统节日，阖家团圆"),
            "元宵节": (2025, 2, 12, "农历正月十五，又称上元节，有赏灯、吃元宵习俗"),
            "龙抬头": (2025, 3, 11, "农历二月初二，民间有剃龙头、祭龙王等习俗"),
            "清明节": (2025, 4, 4, "公历固定日期，兼具自然与人文内涵，扫墓祭祖、踏青插柳"),
            "端午节": (2025, 6, 10, "农历五月初五，纪念屈原，有赛龙舟、吃粽子习俗"),
            "七夕节": (2025, 8, 19, "农历七月初七，又称乞巧节，象征忠贞爱情"),
            "中元节": (2025, 8, 27, "农历七月十五，民间有祭祖、放河灯等习俗"),
            "中秋节": (2025, 9, 29, "农历八月十五，赏月、吃月饼，象征团圆美满"),
            "重阳节": (2025, 10, 11, "农历九月初九，敬老、登高、插茱萸的传统节日"),
            "冬至": (2025, 12, 21, "二十四节气之一，民间有“冬至大如年”之说，吃饺子/汤圆"),
            "腊八节": (2026, 1, 17, "农历十二月初八，喝腊八粥，为春节做准备（跨2025-2026年）"),

            # 公历固定节日
            "元旦": (2025, 1, 1, "公历新年，标志新一年的开始"),
            "妇女节": (2025, 3, 8, "国际妇女节，致敬女性贡献"),
            "植树节": (2025, 3, 12, "倡导植树造林、爱护环境的节日"),
            "劳动节": (2025, 5, 1, "国际劳动节，纪念劳动者的节日"),
            "青年节": (2025, 5, 4, "纪念五四运动，弘扬青年精神"),
            "儿童节": (2025, 6, 1, "国际儿童节，保障儿童权益，庆祝童年"),
            "建党节": (2025, 7, 1, "中国共产党成立纪念日"),
            "建军节": (2025, 8, 1, "中国人民解放军建军纪念日"),
            "国庆节": (2025, 10, 1, "中华人民共和国成立纪念日，举国欢庆"),

            # 重要纪念日
            "教师节": (2025, 9, 10, "致敬教师，弘扬尊师重教传统"),
            "国家公祭日": (2025, 12, 13, "纪念南京大屠杀遇难同胞，铭记历史"),
            "澳门回归纪念日": (2025, 12, 20, "纪念1999年澳门回归祖国")
        }

        # 修复：用(year, month, day, _)解包所有节日（忽略描述字段）
        for holiday, (year, month, day, _) in holidays.items():
            if holiday in query_lower:
                holiday_date = datetime(year, month, day)
                return f"{holiday_date.strftime('%Y-%m-%d')} {calendar.day_name[holiday_date.weekday()]}"

        # 处理直接日期查询（如"2025-10-01星期几"）
        try:
            if "星期几" in query_lower:
                date_str = re.search(r'\d{4}-\d{2}-\d{2}', query).group()
                target_date = datetime.strptime(date_str, '%Y-%m-%d')
                return f"{date_str} {calendar.day_name[target_date.weekday()]}"
        except:
            pass

        return "❌ 无法解析的日期查询，请使用示例格式提问"
    except Exception as e:
        return f"❌ 日期解析失败：{str(e)}"


# ------------------------------
# 3. 智谱风格网络搜索工具
# ------------------------------
class SearchResult:
    def __init__(
            self,
            title: str = "",
            url: str = "",
            summary: str = "",
            site_name: str = "",
            publish_date: str = "",
            favicon: Optional[str] = None
    ):
        self.title = title
        self.url = url
        self.summary = summary
        self.site_name = site_name
        self.publish_date = publish_date
        self.favicon = favicon

    def to_dict(self) -> Dict[str, Any]:
        return {
            "title": self.title,
            "url": self.url,
            "summary": self.summary,
            "site_name": self.site_name,
            "publish_date": self.publish_date,
            "favicon": self.favicon
        }


# ------------------------------
# 4. 高德地图工具
# ------------------------------
async def amap_tools() -> List[Tool]:
    """
    生成高德地图相关工具列表（返回List[Tool]，适配tools.extend()）
    假设amap_mcp_server.get_mcp_client()返回的客户端，其amap_tools()方法已返回Tool列表
    若客户端返回其他格式，需在下方补充转换逻辑（如将接口结果封装为Tool对象）
    """
    try:
        client = await amap_mcp_server.get_amap_mcp_client()

        amap_tool_list = await client.amap_tools()


        if not isinstance(amap_tool_list, list) or (amap_tool_list and not isinstance(amap_tool_list[0], Tool)):
            print("⚠️ 高德地图工具返回非Tool列表，已过滤")
            return []

        print(f"✅ 成功加载 {len(amap_tool_list)} 个高德地图工具")
        return amap_tool_list
    except Exception as e:
        print(f"❌ 加载高德地图工具失败：{str(e)}")
        return []  # 出错时返回空列表，避免影响整体工具加载


# ------------------------------
# 5. 智谱网络搜索工具（保持原有逻辑）
# ------------------------------
@tool("web_search", return_direct=False)
def web_search(**kwargs) -> str:
    """
    【适用场景】
    1. 需要获取**实时/时效性强**的信息（如2024年之后的新闻、当前天气、最新政策）；
    2. 需要验证**不确定的事实**（如某个事件的具体时间、人物职务、数据统计）；
    3. 需要查询**知识库中没有的内容**（如行业最新动态、未存储的技术文档、外部链接信息）；
    4. 需要获取**动态变化的数据**（如股票价格、汇率、热门榜单）。

    【输入格式】
    传入字典格式参数，支持的键（key）包括：
    - search_query: str（必填，搜索关键词，如"2025年4月财经新闻"）
    - count: int（可选，返回结果数量，默认10条，范围1-50）
    - search_engine: str（可选，搜索引擎类型，默认search_std，支持：
      - search_std：基础版（智谱自研，性价比高）
      - search_pro：高级版（智谱自研，召回率/准确率更高）
      - search_pro_sogou：搜狗（覆盖腾讯生态/知乎，垂直领域权威）
      - search_pro_quark：夸克（精准触达垂直内容））
    - search_domain_filter: str（可选，指定域名过滤，如"www.sohu.com"，多个用逗号分隔）
    - search_recency_filter: str（可选，时间范围过滤，默认noLimit，支持：
      - noLimit：无限制
      - 1d：过去1天
      - 7d：过去7天
      - 30d：过去30天
      - 90d：过去90天）
    - content_size: str（可选，摘要字数控制，默认medium，支持low/medium/high）
    - 其他参数：工具会自动忽略无关参数。

    【输出说明】
    返回按相关性排序的搜索结果，每条包含"来源名称"、"发布日期"、"网页标题"、"内容摘要"、"链接"，格式如下：
    1. 来源：搜狐 | 发布日期：2025-05-23
       标题：2025年5月23日财经早资讯
       摘要：1-4月我国对外直接投资575.4亿美元，同比增长7.5%...
       链接：https://www.sohu.com/a/897879632_121123890

    【注意】
    1. 若查询**已存储的固定知识**（如个人简历、项目架构），请调用search_similar_in_knowledge，无需使用此工具；
    2. 不同搜索引擎价格不同（0.01-0.05元/次），优先推荐search_pro平衡效果与成本。
    """
    try:
        search_query = kwargs.get("search_query")
        count = max(1, min(kwargs.get("count", 10), 50))
        search_engine = kwargs.get("search_engine", "search_std")
        search_domain_filter = kwargs.get("search_domain_filter", "")
        search_recency_filter = kwargs.get("search_recency_filter", "noLimit")
        content_size = kwargs.get("content_size", "medium")

        if not search_query or not isinstance(search_query, str):
            return "❌ 搜索失败：缺少必填参数'search_query'（需为字符串类型的搜索关键词）"

        valid_engines = ["search_std", "search_pro", "search_pro_sogou", "search_pro_quark"]
        if search_engine not in valid_engines:
            return f"❌ 搜索失败：无效的search_engine，支持值：{', '.join(valid_engines)}"

        valid_recency = ["noLimit", "1d", "7d", "30d", "90d"]
        if search_recency_filter not in valid_recency:
            return f"❌ 搜索失败：无效的search_recency_filter，支持值：{', '.join(valid_recency)}"

        valid_content_size = ["low", "medium", "high"]
        if content_size not in valid_content_size:
            return f"❌ 搜索失败：无效的content_size，支持值：{', '.join(valid_content_size)}"

        try:
            client = ZhipuAiClient(api_key=settings.Zhipu_search_api_key)
        except Exception as e:
            return f"❌ 客户端初始化失败：{str(e)}（请检查API Key是否正确）"

        print(f"🔍 调用智谱[{search_engine}]搜索：{search_query}（返回{count}条）")
        response = client.web_search.web_search(
            search_engine=search_engine,
            search_query=search_query,
            count=count,
            search_domain_filter=search_domain_filter,
            search_recency_filter=search_recency_filter,
            content_size=content_size
        )

        raw_results = response.get("search_result", [])
        if not raw_results:
            return "❌ 未搜索到相关结果（可能关键词过窄或搜索引擎无数据）"

        formatted_results = []
        for idx, item in enumerate(raw_results[:count], 1):
            result = SearchResult(
                title=item.get("title", "未知标题"),
                url=item.get("link", "无链接"),
                summary=item.get("content", "无摘要")[:200] + "..." if len(item.get("content", "")) > 200 else item.get(
                    "content", "无摘要"),
                site_name=item.get("media", "未知网站"),
                publish_date=item.get("publish_date", "未知日期"),
                favicon=item.get("icon")
            )
            formatted_results.append(
                f"{idx}. 来源：{result.site_name} | 发布日期：{result.publish_date}\n"
                f"   标题：{result.title}\n"
                f"   摘要：{result.summary}\n"
                f"   链接：{result.url}"
            )

        return "\n\n".join(formatted_results)

    except AttributeError as e:
        return f"❌ SDK调用失败：{str(e)}（请检查zai-sdk版本是否最新，建议升级至0.0.4+）"
    except Exception as e:
        return f"❌ 搜索工具运行出错：{str(e)}"


# ------------------------------
# 6. 工具列表整合（异步加载，适配高德工具）
# ------------------------------
async def get_knowledge_tools() -> List[Tool]:
    """
    生成LangChain智能体可用的工具列表（含详细触发说明）
    """
    knowledge = Knowledge()

    tools = [
        # 1. 向知识库添加文本
        Tool(
            name="add_texts_to_knowledge",
            func=lambda texts, metadatas=None, ids=None: knowledge.add_texts(texts, metadatas, ids),
            description="""
            【工具用途】向固定知识库集合中**持久化存储文本知识**。
            【适用场景】
            - 有新的文本内容需要保存（如文档片段、技术笔记、简历信息、项目文档）；
            - 收到用户提供的资料，需要存入知识库供后续查询；
            - 补充或更新知识库中的固定内容。

            【输入参数详解】
            1. texts（必填）：文本列表（即使只有1条，也需用列表包裹），例如：
               ["LangChain支持多向量数据库集成", "Milvus是开源向量数据库"]
            2. metadatas（可选）：元数据列表（与texts长度一致），用于分类和过滤，例如：
               [{"来源":"LangChain文档","类型":"技术笔记"}, {"来源":"Milvus官网","类型":"产品介绍"}]
               若无需元数据，可传入None或空列表；
            3. ids（可选）：文档ID列表（与texts长度一致），仅当知识库开启auto_id=False时有效，否则忽略；
               若无需自定义ID，可传入None。

            【调用示例】
            当用户说"帮我把这两条技术笔记存到知识库：1. LangChain支持Agent；2. Milvus支持IVF_FLAT索引"，
            应调用此工具，输入：
            texts=["LangChain支持Agent", "Milvus支持IVF_FLAT索引"],
            metadatas=[{"类型":"技术笔记"}, {"类型":"技术笔记"}],
            ids=None

            【注意】
            - 文本需是具体的知识内容，而非查询需求；
            - 若用户未提供元数据，可自动添加默认元数据（如{"添加时间":"2025-09-19"}）。
            """,
            return_direct=False
        ),

        # 2. 搜索知识库相似文档
        Tool(
            name="search_similar_in_knowledge",
            func=lambda query, k=5, filter=None: knowledge.search_similar(query, k, filter),
            description="""
            【工具用途】在**已存储的固定知识库**中搜索与查询内容相似的文档。
            【适用场景】
            - 需要获取已存储的知识（如个人简历细节、项目架构、技术栈信息）；
            - 回答用户关于固定知识的问题（如"短链平台如何解决并发扣减？"）；
            - 不需要实时信息，仅需复用已存储内容时。

            【输入参数详解】
            1. query（必填）：查询文本（直接传入用户问题或关键词），例如：
               "AI智能云盘的RAG方案设计"、"短链平台的分库分表策略"
            2. k（可选）：返回结果数量（默认5条，最大不超过20条），例如：3
            3. filter（可选）：过滤条件（按元数据筛选），例如：
               {"来源":"项目文档", "类型":"技术笔记"}
               若无需过滤，可传入None。

            【调用示例】
            当用户问"你参与的短链平台如何支撑日均30亿次访问？"，
            应调用此工具，输入：
            query="短链平台 支撑 日均30亿次访问",
            k=3,
            filter={"项目名称":"短链平台"}

            【注意】
            - 若查询的是**实时/未存储的内容**，应使用web_search工具，而非此工具；
            - 若搜索结果为空，可提示用户"知识库中未找到相关内容，是否需要搜索网络？"。
            """,
            return_direct=False
        ),

        # 3. 删除知识库文档
        Tool(
            name="delete_document_from_knowledge",
            func=lambda document_id: knowledge.delete_document(document_id),
            description="""
            【工具用途】从固定知识库集合中**删除指定ID的单个文档**。
            【适用场景】
            - 知识库中的文档存在错误（如技术参数错误、信息过时）；
            - 需要移除重复存储的文档；
            - 用户明确要求删除某条已存储的内容时。

            【输入参数详解】
            1. document_id（必填）：文档ID（字符串类型），例如：
               "resume_project_ai_disk_001"、"1005"
               （ID是添加文档时返回的inserted_ids，或通过搜索结果获取）

            【调用示例】
            当用户说"帮我删除ID为1005的短链平台文档"，
            应调用此工具，输入：document_id="1005"

            【注意】
            - 必须获取明确的文档ID才能调用，不可凭内容模糊删除；
            - 若用户不知道ID，可先调用search_similar_in_knowledge获取文档ID。
            """,
            return_direct=False
        ),

        # 4. 获取知识库文档数量
        Tool(
            name="get_knowledge_document_count",
            func=lambda: knowledge.get_document_count(),
            description="""
            【工具用途】查询固定知识库集合中**当前的文档总数**。
            【适用场景】
            - 需要了解知识库的规模（如"知识库现在有多少条内容？"）；
            - 验证文档是否成功添加/删除（如"刚才添加的3条笔记存进去了吗？"）；
            - 向用户汇报知识库状态时。

            【输入参数】无（直接调用，无需传入参数）。

            【调用示例】
            当用户问"我的知识库现在有多少条文档？"，
            应调用此工具，无输入参数。

            【注意】
            - 此工具仅返回数量，不返回具体内容；
            - 若需查看具体内容，应调用search_similar_in_knowledge。
            """,
            return_direct=False
        ),

        # 5. 日期解析工具
        Tool(
            name="date_parser",
            func=date_parser,
            description=date_parser.description,
            return_direct=False
        ),



    ]

    # 加载高德地图工具（已修复为返回List[Tool]）
    tool_amap = await amap_tools()
    tools.extend(tool_amap)

    print(f"✅ 工具列表加载完成，共 {len(tools)} 个工具")
    return tools